//
// Created by ZhaoXiaoFei on 2022/10/24.
// Last changed by: ZhaoXiaoFei.
// Last update: 2022/10/24.
//

#ifndef MSCKF_MONO_FILTERING_FLOW_H
#define MSCKF_MONO_FILTERING_FLOW_H

#include <ros/ros.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/Image.h>
#include <nav_msgs/Odometry.h>
#include <deque>
#include <condition_variable>
#include <thread>
#include <chrono>
#include <fstream>
#include <yaml-cpp/yaml.h>
#include "msckf_mono/data_types.h"
#include "msckf_mono/corner_detector.h"
#include "msckf_mono/msckf.h"

class FilteringFlow{

public:
    FilteringFlow(ros::NodeHandle& nh) : nh_(nh), imu_calibrated_(false){
        imu_sub_ = nh_.subscribe("/imu0", 200, &FilteringFlow::ImuCallBack, this);
        image_sub_ = nh_.subscribe("/cam0/image_raw", 20, &FilteringFlow::ImageCallBack, this);
        image_pub_ = nh_.advertise<sensor_msgs::Image>("image", 1);
        odom_pub_ = nh.advertise<nav_msgs::Odometry>("odom", 100);
        LoadParameters();
        corner_track_handler_ = std::make_shared<corner_detector::CornerTrackHandler>(K_, dist_coeffs_,
                                                                                      n_grid_rows_, n_grid_cols_);
        corner_track_handler_->SetRansacParameters(ransac_threshold_);
        my_MSCKF_ = std::make_shared<msckf_mono::MSCKF>();
        msckf_thread_ = std::thread(&FilteringFlow::Process, this);

    }
    ~FilteringFlow() = default;

private:
    void ImuCallBack(const sensor_msgs::ImuConstPtr& imu_msg){
        double cur_imu_time = imu_msg->header.stamp.toSec();

        static bool first_imu_data = false;
        if(!first_imu_data){
            first_imu_data = true;
            //pre_imu_time_ = cur_imu_time;
            imu_stand_time_point_ = cur_imu_time + imu_stand_duration_;
            //return;
        }

        mtx_.lock();
        imu_data_buf_.push_back(imu_msg);
        mtx_.unlock();
        con_.notify_one();
        pre_imu_time_ = cur_imu_time;
    }

    void ImageCallBack(const sensor_msgs::ImageConstPtr & image_msg){
//        double cur_image_time = image_msg->header.stamp.toSec();
//
//        if(!imu_calibrated_){
//            if(pre_imu_time_ > imu_stand_time_point_){
//                mtx_.lock();
//                Initialize_imu();
//                pre_imu_data_.time = imu_data_buf_.back()->header.stamp.toSec();
//                pre_imu_data_.acc.x() = imu_data_buf_.back()->linear_acceleration.x;
//                pre_imu_data_.acc.y() = imu_data_buf_.back()->linear_acceleration.y;
//                pre_imu_data_.acc.z() = imu_data_buf_.back()->linear_acceleration.z;
//                pre_imu_data_.gyr.x() = imu_data_buf_.back()->angular_velocity.x;
//                pre_imu_data_.gyr.y() = imu_data_buf_.back()->angular_velocity.y;
//                pre_imu_data_.gyr.z() = imu_data_buf_.back()->angular_velocity.z;
//
//                imu_data_buf_.clear();
//                mtx_.unlock();
//                imu_calibrated_ = true;
//
//            }
//            return;
//        }

        mtx_.lock();
        image_data_buf_.push_back(image_msg);
        mtx_.unlock();
        con_.notify_one();
    }

    void Process(){
        std::chrono::steady_clock::time_point next_run = std::chrono::steady_clock::now();
        std::chrono::steady_clock::duration duration(std::chrono::milliseconds(50));
        while (true){
            if(!imu_calibrated_){
                if(pre_imu_time_ > imu_stand_time_point_){
                    mtx_.lock();
                    Initialize_imu();
                    pre_imu_data_.time = imu_data_buf_.back()->header.stamp.toSec();
                    pre_imu_data_.acc.x() = imu_data_buf_.back()->linear_acceleration.x;
                    pre_imu_data_.acc.y() = imu_data_buf_.back()->linear_acceleration.y;
                    pre_imu_data_.acc.z() = imu_data_buf_.back()->linear_acceleration.z;
                    pre_imu_data_.gyr.x() = imu_data_buf_.back()->angular_velocity.x;
                    pre_imu_data_.gyr.y() = imu_data_buf_.back()->angular_velocity.y;
                    pre_imu_data_.gyr.z() = imu_data_buf_.back()->angular_velocity.z;

                    imu_data_buf_.clear();
                    mtx_.unlock();
                    my_MSCKF_->Init(camera_, msckf_params_, noise_params_, init_imu_state_);
                    imu_calibrated_ = true;
                }
                continue;
            }
            Work();
            next_run += duration;
            std::this_thread::sleep_until(next_run);
        }
    }

    void Work(){
        std::vector<std::pair<sensor_msgs::ImageConstPtr , std::vector<sensor_msgs::ImuConstPtr>>> measurements;
        std::unique_lock<std::mutex> lk(mtx_);
        con_.wait(lk, [&](){
            return (measurements = GetMeasurements()).size() != 0;
        });
        lk.unlock();
        static std::fstream of("/home/ubuntu/1.txt", std::ios::out);
        for(auto& measurement : measurements){
            double cur_image_time = measurement.first->header.stamp.toSec();
            cv_bridge::CvImagePtr cv_image_ptr;
            try {
                cv_image_ptr = cv_bridge::toCvCopy(measurement.first, sensor_msgs::image_encodings::MONO8);
            }
            catch (cv_bridge::Exception &e) {
                std::cout << "cv_bridge exception: " << e.what() << std::endl;
                continue;
            }
            std::vector<ImuData> imu_before_cur_img;
            float dx = 0, dy = 0, dz = 0, rx = 0, ry = 0, rz = 0;
            double temp_imu_time = -1;
            for (int i = 0; i < measurement.second.size(); i++) {
                double cur_imu_time = measurement.second[i]->header.stamp.toSec();
                if(temp_imu_time < 0){
                    temp_imu_time = cur_imu_time;
                }
                if (cur_imu_time <= cur_image_time) {
                    ImuData current_imu;
                    current_imu.acc.x() = dx = measurement.second[i]->linear_acceleration.x;
                    current_imu.acc.y() = dy = measurement.second[i]->linear_acceleration.y;
                    current_imu.acc.z() = dz = measurement.second[i]->linear_acceleration.z;

                    current_imu.gyr.x() = rx = measurement.second[i]->angular_velocity.x;
                    current_imu.gyr.y() = ry = measurement.second[i]->angular_velocity.y;
                    current_imu.gyr.z() = rz = measurement.second[i]->angular_velocity.z;

                    current_imu.time = cur_imu_time;
                    temp_imu_time = cur_imu_time;
                    imu_before_cur_img.emplace_back(current_imu);
                    of << std::setprecision(13) << "imu1: " << cur_imu_time << std::endl;
                }
                else {
                    double dt_1 = cur_image_time - temp_imu_time;
                    double dt_2 = cur_imu_time - cur_image_time;
                    assert(dt_1 >= 0);
                    assert(dt_2 >= 0);
                    assert(dt_1 + dt_2 > 0);
                    float w1 = dt_2 / (dt_1 + dt_2);
                    float w2 = dt_1 / (dt_1 + dt_2);
                    ImuData current_imu;
                    current_imu.acc.x() = w1 * dx + w2 * measurement.second[i]->linear_acceleration.x;
                    current_imu.acc.y() = w1 * dy + w2 * measurement.second[i]->linear_acceleration.y;
                    current_imu.acc.z() = w1 * dz + w2 * measurement.second[i]->linear_acceleration.z;

                    current_imu.gyr.x() = w1 * rx + w2 * measurement.second[i]->angular_velocity.x;
                    current_imu.gyr.y() = w1 * ry + w2 * measurement.second[i]->angular_velocity.y;
                    current_imu.gyr.z() = w1 * rz + w2 * measurement.second[i]->angular_velocity.z;

                    current_imu.time = cur_image_time;
                    temp_imu_time = cur_image_time;

                    imu_before_cur_img.emplace_back(current_imu);

                    of << std::setprecision(12) << "imu2: " << cur_image_time << std::endl;
                }
            }

            of << std::setprecision(12) << "image: " << cur_image_time << std::endl;

            for (auto &imu_data: imu_before_cur_img){
                my_MSCKF_->Propagate(pre_imu_data_, imu_data);
                Eigen::Vector3f gyro_measurement = R_cam0_imu_ * (imu_data.gyr - init_imu_state_.b_g);
                corner_track_handler_->AddGyrData(gyro_measurement);
                pre_imu_data_ = imu_data;
            }

            corner_track_handler_->SetCurrentImage(cv_image_ptr->image, cur_image_time);

            std::vector<Eigen::Vector2f,Eigen::aligned_allocator<Eigen::Vector2f>> cur_features;
            std::vector<size_t> cur_ids;
            corner_track_handler_->TrackFeatures(cur_features, cur_ids);

            std::vector<Eigen::Vector2f,Eigen::aligned_allocator<Eigen::Vector2f>> new_features;
            std::vector<size_t> new_ids;
            corner_track_handler_->NewFeatures(new_features, new_ids);

            my_MSCKF_->AugmentState(state_id++, cur_image_time);

            my_MSCKF_->Update(cur_features, cur_ids);

            my_MSCKF_->AddFeatures(new_features, new_ids);

            my_MSCKF_->Marginlize();

            my_MSCKF_->pruneEmptyStates();

            cv::Mat res = corner_track_handler_->GetTrackedImage();
            cv::imshow("image", res);
            cv::waitKey(10);

            publish_core(measurement.first->header.stamp);

            if (image_pub_.getNumSubscribers() > 0) {
                cv_bridge::CvImage out_img;
                out_img.header.frame_id = "cam0";
                out_img.header.stamp = measurement.first->header.stamp;
                out_img.encoding = sensor_msgs::image_encodings::TYPE_8UC3;
                out_img.image = corner_track_handler_->GetTrackedImage();
                image_pub_.publish(out_img.toImageMsg());
            }
        }

    }

    std::vector<std::pair<sensor_msgs::ImageConstPtr , std::vector<sensor_msgs::ImuConstPtr>>>
    GetMeasurements(){
        std::vector<std::pair<sensor_msgs::ImageConstPtr , std::vector<sensor_msgs::ImuConstPtr>>> measurement;
        while(true){
            if(imu_data_buf_.empty() || image_data_buf_.empty()){
                return measurement;
            }

            if(imu_data_buf_.back()->header.stamp.toSec() <= image_data_buf_.front()->header.stamp.toSec()){
                return measurement;
            }

            if(imu_data_buf_.front()->header.stamp.toSec() >= image_data_buf_.front()->header.stamp.toSec()){
                image_data_buf_.pop_front();
                continue;
            }

            sensor_msgs::ImageConstPtr img_msg = image_data_buf_.front();
            image_data_buf_.pop_front();
            std::vector<sensor_msgs::ImuConstPtr> imu_msgs;
            while(imu_data_buf_.front()->header.stamp.toSec() <= img_msg->header.stamp.toSec()){
                imu_msgs.push_back(imu_data_buf_.front());
                imu_data_buf_.pop_front();
            }
            imu_msgs.push_back(imu_data_buf_.front());
            if(imu_msgs.empty()){
                std::cout << "no imu data" << std::endl;
            }
            measurement.emplace_back(img_msg, imu_msgs);
        }

        return measurement;
    }

    void Initialize_imu(){
        Eigen::Vector3f acc_mean, gyr_mean;
        acc_mean.setZero();
        gyr_mean.setZero();
        int N = imu_data_buf_.size();
        for(auto data : imu_data_buf_){
            acc_mean += Eigen::Vector3f(data->linear_acceleration.x,
                                        data->linear_acceleration.y,
                                        data->linear_acceleration.z);
            gyr_mean += Eigen::Vector3f(data->angular_velocity.x,
                                        data->angular_velocity.y,
                                        data->angular_velocity.z);
        }
        acc_mean /= N;
        gyr_mean /= N;

        init_imu_state_.p_G_I.setZero();
        init_imu_state_.v_G_I.setZero();
        init_imu_state_.b_g = gyr_mean;
        init_imu_state_.g = Eigen::Vector3f(0, 0, -9.81);

        init_imu_state_.q_I_G = Eigen::Quaternionf::FromTwoVectors(-init_imu_state_.g, acc_mean);
        init_imu_state_.b_a = init_imu_state_.q_I_G * init_imu_state_.g + acc_mean;

        //other
//        Eigen::Vector3f e_1(1, 0, 0);
//        Eigen::Vector3f z_axis = acc_mean / acc_mean.norm();
//        Eigen::Vector3f x_axis = e_1 - (z_axis.transpose() * e_1) * z_axis; x_axis /= x_axis.norm();
//        Eigen::Vector3f y_axis = z_axis.cross(x_axis);
//        Eigen::Matrix3f R_I_G;
//        R_I_G.block<3, 1>(0, 0) = x_axis;
//        R_I_G.block<3, 1>(0, 1) = y_axis;
//        R_I_G.block<3, 1>(0, 2) = z_axis;
//        std::cout << Eigen::Quaternionf(R_I_G) << std::endl;

        const auto q = init_imu_state_.q_I_G;
        std::cout << "\nInitial IMU State" <<
                                              "\n--p_G_I " << init_imu_state_.p_G_I.transpose() <<
                                              "\n--q_I_G " << q.w() << ", " << q.x() << ", " << q.y() << ", " << q.z() <<
                                              "\n--v_G_I " << init_imu_state_.v_G_I.transpose() <<
                                              "\n--b_a " << init_imu_state_.b_a.transpose() <<
                                              "\n--b_g " << init_imu_state_.b_g.transpose() <<
                                              "\n--g " << init_imu_state_.g.transpose() << std::endl;

    }

    void LoadParameters(){
        YAML::Node node = YAML::LoadFile("/home/ubuntu/WorkSpace/VSlam/msckf_ws/src/msckf_mono/config/config.yaml");
        std::string kalibr_camera = "cam0";

        camera_model_ = node[kalibr_camera]["camera_model"].as<std::string>();

        K_ = cv::Mat::eye(3, 3, CV_32F);
        std::vector<float> intrinsics(4);
        intrinsics = node[kalibr_camera]["intrinsics"].as<std::vector<float>>();
        K_.at<float>(0, 0) = intrinsics[0];
        K_.at<float>(1, 1) = intrinsics[1];
        K_.at<float>(0, 2) = intrinsics[2];
        K_.at<float>(1, 2) = intrinsics[3];

        distortion_model_ = node[kalibr_camera]["distortion_model"].as<std::string>();
        std::vector<float> distortion_coeffs(4);
        distortion_coeffs = node[kalibr_camera]["distortion_coeffs"].as<std::vector<float>>();
        dist_coeffs_ = cv::Mat::zeros(distortion_coeffs.size(), 1, CV_32F);
        dist_coeffs_.at<float>(0) = distortion_coeffs[0];
        dist_coeffs_.at<float>(1) = distortion_coeffs[1];
        dist_coeffs_.at<float>(2) = distortion_coeffs[2];
        dist_coeffs_.at<float>(3) = distortion_coeffs[3];


        Eigen::Matrix4f T_imu_cam;
        for(int i = 0; i < node[kalibr_camera]["T_cam_imu"].size(); i++){
            for(int j = 0; j < node[kalibr_camera]["T_cam_imu"][i].size(); j++){
                T_imu_cam(i, j) = node[kalibr_camera]["T_cam_imu"][i][j].as<float>();
            }
        }

        R_imu_cam0_ = T_imu_cam.block<3, 3>(0, 0);
        p_imu_cam0_ = T_imu_cam.block<3, 1>(0, 3);


        R_cam0_imu_ = R_imu_cam0_.transpose();
        p_cam0_imu_ = R_cam0_imu_ * (-1. * p_imu_cam0_);

        // setup camera parameters
        camera_.f_u = intrinsics[0];
        camera_.f_v = intrinsics[1];
        camera_.c_u = intrinsics[2];
        camera_.c_v = intrinsics[3];

        camera_.q_C_I = Eigen::Quaternionf(R_imu_cam0_).inverse();
        camera_.p_I_C = p_imu_cam0_;

        n_grid_rows_ = node["n_grid_rows"].as<int>();
        n_grid_cols_ = node["n_grid_cols"].as<int>();

        ransac_threshold_ = node["ransac_threshold"].as<float>();

        // MSCKF Parameters
        float feature_cov;
        feature_cov = node["feature_covariance"].as<float>();

        Eigen::Matrix<float, 12, 1> Q_imu_vars;
        float w_var, dbg_var, a_var, dba_var;
        w_var = node["imu_vars"]["w_var"].as<float>();
        dbg_var = node["imu_vars"]["dbg_var"].as<float>();
        a_var = node["imu_vars"]["a_var"].as<float>();
        dba_var = node["imu_vars"]["dba_var"].as<float>();

        Q_imu_vars << w_var, w_var, w_var,
                        dbg_var, dbg_var, dbg_var,
                        a_var, a_var, a_var,
                        dba_var, dba_var, dba_var;

        Eigen::Matrix<float, 15, 1> IMUCovar_vars;
        float q_var_init, bg_var_init, v_var_init, ba_var_init, p_var_init;
        p_var_init = node["imu_covars"]["p_var_init"].as<float>();
        q_var_init = node["imu_covars"]["q_var_init"].as<float>();
        v_var_init = node["imu_covars"]["v_var_init"].as<float>();
        ba_var_init = node["imu_covars"]["ba_var_init"].as<float>();
        bg_var_init = node["imu_covars"]["bg_var_init"].as<float>();
        IMUCovar_vars << q_var_init, q_var_init, q_var_init,
                            bg_var_init, bg_var_init, bg_var_init,
                            v_var_init, v_var_init, v_var_init,
                            ba_var_init, ba_var_init, ba_var_init,
                            p_var_init, p_var_init, p_var_init;

        // Setup noise parameters
        noise_params_.init_imu_covar = IMUCovar_vars.asDiagonal();
        noise_params_.Q_imu = Q_imu_vars.asDiagonal();
        noise_params_.u_var = pow(feature_cov / camera_.f_u, 2);
        noise_params_.v_var = pow(feature_cov / camera_.f_v, 2);

        msckf_params_.max_gn_cost_norm = node["max_gn_cost_norm"].as<float>();
        msckf_params_.max_gn_cost_norm = pow(msckf_params_.max_gn_cost_norm / camera_.f_u, 2);
        msckf_params_.translation_threshold = node["translation_threshold"].as<float>();
        msckf_params_.min_rcond = node["min_rcond"].as<float>();
        msckf_params_.redundancy_angle_thresh = node["keyframe_transl_dist"].as<float>();
        msckf_params_.redundancy_distance_thresh = node["keyframe_rot_dist"].as<float>();
        msckf_params_.max_track_length = node["max_track_length"].as<int>();
        msckf_params_.min_track_length = node["min_track_length"].as<int>();
        msckf_params_.max_cam_states = node["max_cam_states"].as<int>();

        imu_stand_duration_ = node["stand_still_time"].as<double>();

        std::cout << "ransac_threshold_: " << ransac_threshold_ << std::endl;
        std::cout << "Loaded " << kalibr_camera;
        std::cout << "-Intrinsics " << intrinsics[0] << ", "
                                       << intrinsics[1] << ", "
                                       << intrinsics[2] << ", "
                                       << intrinsics[3] << std::endl;;
        std::cout << "-Distortion " << distortion_coeffs[0] << ", "
                                       << distortion_coeffs[1] << ", "
                                       << distortion_coeffs[2] << ", "
                                       << distortion_coeffs[3] << std::endl;;
        const auto q_CI = camera_.q_C_I;
        std::cout << "-q_CI \n" << q_CI.x() << "," << q_CI.y() << "," << q_CI.z() << "," << q_CI.w() << std::endl;;
        std::cout << "-p_I_C \n" << camera_.p_I_C.transpose() << std::endl;
    }

    void publish_core(const ros::Time &publish_time) {
        auto imu_state = my_MSCKF_->GetImuState();

        nav_msgs::Odometry odom;
        odom.header.stamp = publish_time;
        odom.header.frame_id = "map";
        odom.twist.twist.linear.x = imu_state.v_G_I[0];
        odom.twist.twist.linear.y = imu_state.v_G_I[1];
        odom.twist.twist.linear.z = imu_state.v_G_I[2];

        odom.pose.pose.position.x = imu_state.p_G_I[0];
        odom.pose.pose.position.y = imu_state.p_G_I[1];
        odom.pose.pose.position.z = imu_state.p_G_I[2];
        Eigen::Quaternionf q_out = imu_state.q_I_G.inverse();
        odom.pose.pose.orientation.w = q_out.w();
        odom.pose.pose.orientation.x = q_out.x();
        odom.pose.pose.orientation.y = q_out.y();
        odom.pose.pose.orientation.z = q_out.z();

        odom_pub_.publish(odom);
    }

private:
    ros::NodeHandle nh_;
    ros::Subscriber imu_sub_;
    ros::Subscriber image_sub_;
    ros::Publisher image_pub_;
    ros::Publisher odom_pub_;

    std::mutex mtx_;
    std::condition_variable con_;
    std::deque<sensor_msgs::ImuConstPtr> imu_data_buf_;
    std::deque<sensor_msgs::ImageConstPtr> image_data_buf_;

    std::thread msckf_thread_;

    std::atomic<bool> imu_calibrated_;
    double imu_stand_duration_;
    double imu_stand_time_point_;
    double pre_imu_time_ = 0.0;

    ImuData pre_imu_data_;
    ImuState init_imu_state_;

    std::shared_ptr<corner_detector::CornerTrackHandler> corner_track_handler_;
    std::shared_ptr<msckf_mono::MSCKF> my_MSCKF_;

    std::string camera_model_, distortion_model_;
    cv::Mat K_, K_inv_, dist_coeffs_;

    Eigen::Matrix3f R_imu_cam0_; Eigen::Vector3f p_imu_cam0_;
    Eigen::Matrix3f R_cam0_imu_; Eigen::Vector3f p_cam0_imu_;

    int n_grid_rows_, n_grid_cols_;
    Camera camera_;
    NoiseParams noise_params_;
    MSCKFParams msckf_params_;

    int state_id = 0;
    float ransac_threshold_;
};

#endif //MSCKF_MONO_FILTERING_FLOW_H
